#!/usr/bin/env python3
import inspect
from functools import wraps

class ToolRegistry:
    def __init__(self):
        # self.tools = []
        self.tools = {}

    def tool(self, func):
        """
        A decorator to annotate tool functions, extracting their metadata such as name, description, and arguments.
        """

        # Use wraps to preserve the function's metadata
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        # register the tool function
        # Extract the function name
        wrapper.name = func.__name__

        # Extract the docstring as the description
        wrapper.description = func.__doc__.strip() if func.__doc__ else ""

        # Extract argument names, types, and defaults using inspect.signature
        signature = inspect.signature(func)
        wrapper.args = {
            param.name: {
                "title": param.name.capitalize(),
                "type": param.annotation.__name__ if param.annotation != inspect._empty else "any",
            }
            for param in signature.parameters.values()
        }
        arg_str = ', '.join(f'{k}: {v["type"]}' for k, v in wrapper.args.items())
        wrapper.meta = f"{wrapper.name}({arg_str}) -> str - {wrapper.description}"

        self.tools[wrapper.name] = wrapper
        # self.tools.append(wrapper)

        return wrapper

    def render_text_description(self):
        """
        Generate a textual description of tools' metadata.

        Returns:
            str: A string containing the metadata of all tools in a readable format.
        """
        descriptions = []

        for tool in self.tools.values():
            descriptions.append(tool.meta)

        # Join all tool descriptions with newlines
        return "\n".join(descriptions)

    def render_json_description(self):
        """
        Generate a JSON description of tools' metadata.

        Returns:
            str: A JSON string containing the metadata of all tools in a structured format.
        """
        tools_info = {}
        for name, tool in self.tools.items():
            tools_info[name] = {
                "name": tool.name,
                "description": tool.description,
                "arguments": tool.args
            }
        return json.dumps(tools_info, indent=4)

def tool(func):
    """
    A decorator to annotate tool functions, extracting their metadata such as name, description, and arguments.
    """
    # Use wraps to preserve the function's metadata
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)

    # Extract the function name
    wrapper.name = func.__name__

    # Extract the docstring as the description
    wrapper.description = func.__doc__.strip() if func.__doc__ else ""

    # Extract argument names, types, and defaults using inspect.signature
    signature = inspect.signature(func)
    wrapper.args = {
        param.name: {
            "title": param.name.capitalize(),
            "type": param.annotation.__name__ if param.annotation != inspect._empty else "any",
        }
        for param in signature.parameters.values()
    }
    arg_str = ', '.join(f'{k}: {v["type"]}' for k, v in wrapper.args.items())
    wrapper.meta = f"{wrapper.name}({arg_str}) -> str - {wrapper.description}"

    return wrapper

def render_text_description(tools):
    """
    Generate a textual description of tools' metadata.

    Args:
        tools (list): A list of tools, where each tool has `name`, `description`, and `args` attributes.

    Returns:
        str: A string containing the metadata of all tools in a readable format.
    """
    descriptions = []

    for tool in tools:
        descriptions.append(tool.meta)

    # Join all tool descriptions with newlines
    return "\n".join(descriptions)


if __name__=="__main__":

    # Example 1: Define a tool using the @tool decorator
    @tool
    def search(query: str) -> str:
        """Look up things online."""
        return "Some results"

    # Example 2: Define another tool with more parameters
    @tool
    def calculate(a: int, b: int) -> int:
        """Perform a basic addition of two numbers."""
        return a + b

    # Print metadata for the tool
    print(search.name)             # "search"
    print(search.description)      # "Look up things online."
    print(search.args)             # {'query': {'title': 'Query', 'type': 'string'}}
    print(search.meta)             # search(query: str) -> str - Look up things online.

    print(calculate.name)          # "calculate"
    print(calculate.description)   # "Perform a basic addition of two numbers."
    print(calculate.args)          # {'a': {'title': 'A', 'type': 'int'}, 'b': {'title': 'B', 'type': 'int'}}
    print(calculate.meta)          # calculate(a: int, b: int) -> str - Perform a basic addition of two numbers.

    # List of tools
    tools = [search, calculate]

    # Render tool descriptions
    rendered_tools = render_text_description(tools)
    print(rendered_tools)
